import cv2
import myfuncs as mf
from ColorFilter import ColorFilter
import numpy as np
import matplotlib.pyplot as plt
import pytesseract as ptr
import os
import json

from Conditions import Conditions
from Kernels import Kernels
from MatchDigit import MatchDigit


class MatchConsumeType:
	ACCOUNT_TYPES_PATH = '.\\templates\\AccountTypes.json'
	TEMPLATE_INDEX_PATH = {
		# 'DATETIME_64x64_Templates': '.\\templates\\DATETIME_64x64_Templates.json'
	}
	originImage = None
	grayImage = None
	binaryImage = None
	imageROI = None
	matchResults = []
	kernels = Kernels()

	def __init__(self, imagePath, width):
		ptr.pytesseract.tesseract_cmd = r"D:\Softwares\Tesseract-OCR\tesseract.exe"
		self.imageInitialize(imagePath, width)

	def imageInitialize(self, imagePath, width):
		originImage = cv2.imread(imagePath)
		originImage = mf.resize(originImage, width=width)
		self.grayImage = cv2.cvtColor(originImage, cv2.COLOR_BGR2GRAY)
		colorFilter = ColorFilter(originImage)
		colorFilter.addColor('black', (0, 0, 0), (180, 50, 222))
		colorFilter.addColor('color', (0, 50, 0), (180, 255, 255))
		blackArea = colorFilter.filterColor('black')[1]
		colorArea = colorFilter.filterColor('color')[1]
		# mf.showImages({
		# 	'origin': originImage,
		# 	'blue': colorArea,
		# 	'black': blackArea,
		# 	'bwr': cv2.bitwise_or(blackArea.copy(), colorArea.copy())
		# })
		self.binaryImage = cv2.bitwise_or(blackArea.copy(), colorArea.copy())
		self.originImage = originImage

	def calcROI(self, operatedArea=None):
		imageROI = self.binaryImage.copy()
		if operatedArea is not None:
			operatedArea = mf.resize(operatedArea, width=self.originImage.shape[1])
			imageROI = cv2.bitwise_or(imageROI, imageROI, mask=operatedArea.copy())
		# mf.showImage('imageROI', imageROI)
		imageROI = cv2.morphologyEx(imageROI.copy(), cv2.MORPH_CLOSE, self.kernels.quickKernel('7r2'))
		imageROI = cv2.morphologyEx(imageROI.copy(), cv2.MORPH_DILATE, self.kernels.quickKernel('3r3'))
		imageROI = cv2.morphologyEx(imageROI.copy(), cv2.MORPH_DILATE, self.kernels.quickKernel('14r2'))
		self.imageROI = imageROI
		# mf.showImage('imageROI', imageROI)

	@staticmethod
	def findContours(roi):
		empty = np.zeros([roi.shape[0], roi.shape[1]], np.uint8)
		contours, hierarchy = cv2.findContours(roi.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
		contoursImage = cv2.drawContours(empty, contours, -1, (200, 20, 20))
		return contours, contoursImage

	def findRectROI(self, roi):
		contours, _ = self.findContours(roi)
		rectParams = []
		rectROIs = []
		rp = []
		x, y, w, h = 0, 0, 0, 0
		for (i, contour) in enumerate(contours):
			(x, y, w, h) = cv2.boundingRect(contour)
			ratio = h / w
			ratioConditions = Conditions()
			ratioConditions.addCondition('not square', not abs(ratio - 1) < 0.2)
			# ratioConditions.addCondition('not long', (not abs(1 / ratio - 100) < 50 and not w > self.originImage.shape[1] * 0.5))
			ratioConditions.addCondition('not big', (not abs(1 / ratio - 2) < 1))
			if ratioConditions.isAllTrue():
				rectParams.append((x, y, w, h))
		for (i, rectParam) in enumerate(rectParams):
			if i == 0:
				continue
			x, y, w, h = rectParam
			px, py, pw, ph = rectParams[i-1]
			rectROI = {}
			relativeConditions = Conditions()
			relativeConditions.addCondition('height > prevHeight', h > ph)
			relativeConditions.addCondition('x close to prevX', abs(x - px) < 3)
			if relativeConditions.isAllTrue():
				rp.append((x, y, w, h))
				rp.append((px, py, pw, ph))
				rectROI['accPosition'] = (x, y, w, h)
				rectROI['datPosition'] = (px, py, pw, ph)
				rectROI['accountROI'] = self.binaryImage[y - 5:y + h + 5, x - 5:x + w + 5]
				rectROI['datetimeROI'] = self.binaryImage[py - 5:py + ph + 5, px - 5:px + pw + 5]
				rectROIs.append(rectROI)
			# 	plt.subplot(121), plt.title('accountROI'), plt.imshow(rectROI['accountROI'])
			# 	plt.subplot(122), plt.title('datetimeROI'), plt.imshow(rectROI['datetimeROI'])
			# plt.show()
		rects = mf.drawRects(rectParams, roi.shape)
		filteredRects = mf.drawRects(rp, roi.shape)
		# mf.showImages({
		# 	'binary': self.binaryImage,
		# 	'rects': rects,
		# 	'filteredRects': filteredRects
		# })
		return rectROIs

	def matchAccountType(self, accountROIs):
		results = []
		accountTypes = mf.readJsonFile2Python(self.ACCOUNT_TYPES_PATH)
		for (i, groupROI_info) in enumerate(accountROIs):
			groupROI = groupROI_info['accountROI']
			x, y, w, h = groupROI_info['accPosition']
			groupResult = {}
			# mf.showImage('groupROI', groupROI)
			text = ptr.image_to_string(groupROI, lang='chi_sim')
			text = text.replace(' ', '')

			for content in accountTypes.items():
				keywords = content[1]['keywords']
				accountType = content[1]['name']
				# print("accType: {}, keywords: {}".format(accountType, keywords))
				for keyword in keywords:
					if keyword in text:
						groupResult = {
							'id': i,
							'position': (x, y),
							'size': (w, h),
							'content': accountType
						}
						break
			if groupResult == {}:
				groupResult = {
					'id': i,
					'position': (x, y),
					'size': (w, h),
					'content': '未分类'
				}
			print(groupResult)
			results.append(groupResult)
		return results

	def matchDatetime(self, datetimeROIs, templates):
		results = []
		templateIndexes = mf.readJsonFile2Python(self.TEMPLATE_INDEX_PATH['DATETIME_64x64_Templates'])
		for (i, groupROI_info) in enumerate(datetimeROIs):
			groupROI = groupROI_info['datetimeROI']
			groupROI = mf.resize(groupROI.copy(), height=templates[(0, 0)].shape[0])
			# mf.showImage('groupROI', groupROI)
			gx, gy, gw, gh = groupROI_info['datPosition']
			contours, hierarchy = cv2.findContours(groupROI.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
			text = ''
			for contour in contours:
				x, y, w, h = cv2.boundingRect(contour)
				charROI = groupROI[y:y+h, x:x+w]
				charROI = cv2.threshold(charROI.copy(), 0, 255, cv2.THRESH_BINARY)[1]
				scores = []
				tags = []
				for (templateIndex, template_info) in enumerate(templates.items()):
					tag, template = template_info
					tmpCnt, _ = cv2.findContours(template.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
					if len(tmpCnt) == 0:
						continue
					tx, ty, tw, th = cv2.boundingRect(tmpCnt[0])
					template = template[ty:ty+th, tx:tx+tw]
					template = cv2.morphologyEx(template.copy(), cv2.MORPH_DILATE, self.kernels.quickKernel('5r5'))
					template = mf.resize(template.copy(), width=charROI.shape[1])
					result = cv2.matchTemplate(charROI, template, cv2.TM_CCOEFF)
					(_, score, _, _) = cv2.minMaxLoc(result)
					scores.append(score)
					tags.append(tag)
				# mf.showImage('matched', templates[tags[np.argmax(scores)]])
				# plt.subplot(121), plt.imshow(templates[tags[np.argmax(scores)]])
				# plt.subplot(122), plt.imshow(charROI)
				# plt.show()
				(c, r) = tags[np.argmax(scores)]
				# print(templateIndexes[r][c])
				text += templateIndexes[r][c]
			text = text[::-1][:-2]
			text = text[0:-2] + ':' + text[-2:]
			# print(text)
			groupResult = {
				'id': i,
				'position': (gx, gy),
				'size': (gw, gh),
				'content': text
			}
			results.append(groupResult)
		return results

	def bindTemplateIndexes(self, tag, path):
		self.TEMPLATE_INDEX_PATH[tag] = path


if __name__ == '__main__':
	mct = MatchConsumeType('operand1.jpg', width=360, singleTemplateSize=(12, 12))
	mct.calcROI()
